#include <iostream>
#include <vector>
#include <map>
#include <set>
#include <algorithm>

#include <cmath>

#include <matrix.h>
#include <cluster.h>

#include <sds/sfile.h>
#include <sds/sdataset.h>
#include <sds/sdataview.h>

#include <lemon/smart_graph.h>
#include <lemon/unionfind.h>
#include <lemon/random.h>
#include <lemon/matching.h>

using namespace std;
using namespace sds;
using namespace lemon;

#if !defined(INIT_NORMAL) && !defined(INIT_ADV)
#define INIT_ADV
#endif

#if !defined(JOIN_NORMAL) && !defined(JOIN_MATCHING)
#define JOIN_MATCHING
#endif


// parameters

int smallClusterSize = 230; //256
double smallDistance = 1.44; //1.64

double maxDistance = 128.0;
double maxDistanceLast = 96.0;

int stepNum = 10;

double initialCluster = 160;

// global variables


double invMaxDistance = inv_map_euclid(maxDistance);
double invMaxDistanceLast = inv_map_euclid(maxDistanceLast);

class IntMap : public std::vector<int> {
public:
  typedef int Key;

  IntMap(int size, int value) : std::vector<int>(size, value) {}
  void set(int key, int value) { (*this)[key] = value;}
};

void classify(const SDataView<Point>& src,
	      std::vector<Point>& centers,
	      std::vector<Cluster>& clusters) {
  clusters.clear();
  clusters.resize(centers.size());
  SDataView<Point>::Data::const_iterator it;
  for (it = src.data().begin(); it != src.data().end(); ++it) {
    if (it->zero()) continue;
    int minPos = 0;
    double minVal = mon_distance_euclid(*it, centers[0]);
    for (int i = 1; i < (int)centers.size(); ++i) {
      double val = mon_distance_euclid(*it, centers[i]);
      if (val < minVal) {
	minVal = val;
	minPos = i;
      }
    }
    if (mon_distance_euclid(*it, centers[minPos]) < invMaxDistance) {
      clusters[minPos] += *it;
    } else {
      centers.push_back(*it);
      clusters.push_back(Cluster());
      clusters.back() += *it;
    }
  }
}

#if defined(INIT_NORMAL)
void start(const SDataView<Point>& src, 
	   std::vector<Point>& centers) {
  Cluster info;
  SDataView<Point>::Data::const_iterator it;
  for (it = src.data().begin(); it != src.data().end(); ++it) {
    if (it->zero()) continue;
    info += *it;
  }
  Vector expected(CH);
  Matrix covariance(CH, CH);
  for (int i = 0; i < CH; ++i) {
    for (int j = 0; j < CH; ++j) {
      covariance(i, j) = info.covariance(i, j);
    }
    expected(i) = info.expected(i);
  }
  Matrix cholesky = covariance.cholesky();
  for (int k = 0; k < initialCluster; ++k) {
    Vector v(CH);
    for (int i = 0; i < CH; ++i) {
      v(i) = lemon::rnd.gauss();
    }
    //    v = exp((double)k / (double)initialCluster) * (v * cholesky) + expected;
    v = v * cholesky + expected;
    Point p;
    for (int i = 0; i < CH; ++i) {
      p[i] = v(i);
    }
    centers.push_back(p);
  }
}

#elif defined(INIT_ADV)

void start(const SDataView<Point>& src, 
	   std::vector<Point>& centers) {

  double dss;
  SDataView<double> ds(src.getSize());

  {
    int n = lemon::rnd[src.data().size()];
    while (src.data()[n].zero()) {
      n = lemon::rnd[src.data().size()];
    }
    Point cntr = src.data()[n];
    centers.push_back(cntr);

    dss = 0.0;
    for (int i = 0; i < src.data().size(); ++i) {
      if (src.data()[i].zero()) continue;
      ds.data()[i] = mon_distance_euclid(src.data()[i], cntr);
      dss += ds.data()[i];
    }
  }

  for (int k = 1; k < initialCluster; ++k) {
    std::cerr << k << std::endl;
    double r = lemon::rnd(dss);

    int n;
    for (n = 0; ; ++n) {
      if (src.data()[n].zero()) continue;
      if (r < ds.data()[n]) break;
      r -= ds.data()[n];
    }

    Point cntr = src.data()[n];
    centers.push_back(cntr);

    dss = 0.0;
    for (int i = 0; i < src.data().size(); ++i) {
      if (src.data()[i].zero()) continue;
      ds.data()[i] = 
	std::min(ds.data()[i], mon_distance_euclid(src.data()[i], cntr));
      dss += ds.data()[i];
    }
  }
}

#endif

void Read(std::string filename, SDataView<Point>& src, vector<SInteger>& size) {
  SFile file(filename);

  SDataView<unsigned char> views[CH];
  
  for (int i = 0; i < CH; ++i) { 
    SDataset dataset = file.getDataset(i);
    views[i] = dataset.getDataView<unsigned char>();
  }

  size = views[0].getSize();

  src = SDataView<Point>(size);

  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      Point item;
      for (int k = 0; k < CH; ++k) {
	item[k] = (double)views[k][i][j];
      }
      src[i][j] = item;
    }
  } 
}

void Write(std::string targetName, const SDataView<Point>& src,
	   const std::vector<Point>& centers, const vector<SInteger>& size) {
  SFile target(targetName, SFile::CREATE);
  SDataset targetDataset = target.createDataset("clustered", UINT8, size);

  SDataView<unsigned char> targetView(size);

  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      targetView[i][j] = 0;
      if (((const Point&)src[i][j]).zero()) continue;
      int minPos = 0;
      double minVal = mon_distance_euclid(src[i][j], centers[0]);
      for (int l = 1; l < (int)centers.size(); ++l) {
	double val = mon_distance_euclid(src[i][j], centers[l]);
	if (val < minVal) {
	  minVal = val;
	  minPos = l;
	}
      }
      if (mon_distance_euclid(src[i][j], centers[minPos]) 
          < invMaxDistanceLast) {
	targetView[i][j] = (unsigned char)minPos + 1;
      }
    }
  }
  targetDataset.setDataView<unsigned char>(targetView);  
}

void removeSmallClusters(std::vector<Cluster>& infos) {
  for (int i = 0; i < (int)infos.size(); ++i) {
    if (infos[i].size() < smallClusterSize) {
      infos[i] = infos.back();
      infos.pop_back();
      --i;
    }
  }
}

#if defined(JOIN_NORMAL)

void joinClusters(std::vector<Cluster>& infos, 
		  std::vector<Point>& centers) {
  centers.clear();
  IntMap unionfindMap(infos.size(), -1);
  UnionFind<IntMap > unionfind(unionfindMap);
  
  for (int i = 0; i < (int)infos.size(); ++i) {
    unionfind.insert(i);
  }

  for (int i = 0; i < (int)infos.size(); ++i) {
    for (int j = i + 1; j < (int)infos.size(); ++j) {
      if (distance_bat(infos[i], infos[j]) < smallDistance) {
	unionfind.join(i, j);
      }
    }
  }
  std::map<int, Cluster> nc;
  std::map<int, int> nn;
  for (int i = 0; i < (int)infos.size(); ++i) {
    nc[unionfind.find(i)] += infos[i];
    nn[unionfind.find(i)] += 1;
  }
  std::map<int, Cluster>::iterator it;
  for (it = nc.begin(); it != nc.end(); ++it) {
    centers.push_back(it->second.center());
    std::cout << nn[it->first] << ' ';
  }
  std::cout << std::endl;
}

#elif defined(JOIN_MATCHING)

void joinClusters(std::vector<Cluster>& infos, 
		  std::vector<Point>& centers) {

  lemon::SmartGraph ugraph;
  std::vector<lemon::SmartGraph::Node> nodes(infos.size());
  lemon::SmartGraph::EdgeMap<double> weight(ugraph);

  centers.clear();

  for (int i = 0; i < (int)infos.size(); ++i) {
    nodes[i] = ugraph.addNode();
  }

  for (int i = 0; i < (int)infos.size(); ++i) {
    for (int j = i + 1; j < (int)infos.size(); ++j) {
      double dist = distance_bat(infos[i], infos[j]);
      if (dist < smallDistance) {
	lemon::SmartGraph::Edge e = ugraph.addEdge(nodes[i], nodes[j]);
	weight[e] = smallDistance - dist;
      }
    }
  }
  lemon::MaxWeightedMatching<lemon::SmartGraph,
    lemon::SmartGraph::EdgeMap<double> > mwm(ugraph, weight);
  mwm.run();

  for (int i = 0; i < (int)infos.size(); ++i) {
    if (mwm.matching(nodes[i]) != INVALID) {
      int mi = ugraph.id(mwm.mate(nodes[i])); 
      if (mi < i) continue;
      Cluster merge = infos[i] + infos[mi];
      centers.push_back(merge.center());
    } else {
      centers.push_back(infos[i].center());
    }
  }
  std::cout << std::endl;
}


#endif


int main(int argc, const char *argv[]) {
  lemon::rnd.seed(time(0));

  assert(argc == 3);
  SDataView<Point> src;
  vector<SInteger> size;
  Read(argv[1], src, size);
  std::vector<Point> centers;
  start(src, centers);
  for (int step = 0; step < stepNum; ++step) {
    std::cout << (double)step / stepNum << std::endl;
    std::vector<Cluster> clusters;
    classify(src, centers, clusters);
    if (step > 0) { 
      removeSmallClusters(clusters);
      joinClusters(clusters, centers);
    } else {
      for (int i = 0; i < clusters.size(); ++i) {
        centers[i] = clusters[i].center();
      } 
    }
    cout << "Num: " << centers.size() << endl;
  }
  Write(argv[2], src, centers, size);

  return 0;
}
